import sqlite3
import json
from datetime import datetime
import sqlite3

class InvestmentHistoryDB:
    def __init__(self, db_path='investment_history.db'):
        """初始化数据库连接"""
        self.db_path = db_path
        self._create_table()
        
    def _create_table(self):
        """创建投资历史表"""
        try:
            with sqlite3.connect(self.db_path) as conn:
                cursor = conn.cursor()
                # 创建投资历史表，如果不存在
                cursor.execute('''
                    CREATE TABLE IF NOT EXISTS investment_history (
                        id INTEGER PRIMARY KEY AUTOINCREMENT,
                        user_id INTEGER NOT NULL,
                        date TEXT NOT NULL,
                        cny REAL NOT NULL,
                        usd REAL NOT NULL DEFAULT 0,
                        exchange_rate REAL NOT NULL DEFAULT 0,
                        category TEXT DEFAULT '货币基金',
                        deposit_start_date TEXT,
                        deposit_duration INTEGER,
                        deposit_rate REAL,
                        created_at TEXT DEFAULT CURRENT_TIMESTAMP,
                        updated_at TEXT DEFAULT CURRENT_TIMESTAMP
                    )
                ''')
                # 如果表已存在但没有某些字段，添加这些字段
                cursor.execute("PRAGMA table_info(investment_history)")
                columns = [column[1] for column in cursor.fetchall()]
                
                if 'category' not in columns:
                    cursor.execute("ALTER TABLE investment_history ADD COLUMN category TEXT DEFAULT '货币基金'")
                
                if 'deposit_start_date' not in columns:
                    cursor.execute("ALTER TABLE investment_history ADD COLUMN deposit_start_date TEXT")
                
                if 'deposit_duration' not in columns:
                    cursor.execute("ALTER TABLE investment_history ADD COLUMN deposit_duration INTEGER")
                
                if 'deposit_rate' not in columns:
                    cursor.execute("ALTER TABLE investment_history ADD COLUMN deposit_rate REAL")
                
                # 将usd和exchange_rate设置为可为空且有默认值
                if 'usd' in columns:
                    # 检查当前usd列是否为NOT NULL
                    cursor.execute("SELECT sql FROM sqlite_master WHERE type='table' AND name='investment_history'")
                    create_table_sql = cursor.fetchone()[0]
                    if 'usd REAL NOT NULL' in create_table_sql or 'usd REAL' in create_table_sql:
                        # 为了修改字段属性，我们需要重新创建表
                        # 1. 创建临时表
                        cursor.execute('''
                            CREATE TABLE IF NOT EXISTS investment_history_temp (
                                id INTEGER PRIMARY KEY AUTOINCREMENT,
                                user_id INTEGER NOT NULL,
                                date TEXT NOT NULL,
                                cny REAL NOT NULL,
                                usd REAL DEFAULT 0,
                                exchange_rate REAL DEFAULT 0,
                                category TEXT DEFAULT '货币基金',
                                deposit_start_date TEXT,
                                deposit_duration INTEGER,
                                deposit_rate REAL,
                                created_at TEXT DEFAULT CURRENT_TIMESTAMP,
                                updated_at TEXT DEFAULT CURRENT_TIMESTAMP
                            )
                        ''')
                        # 2. 复制数据
                        cursor.execute('''
                            INSERT INTO investment_history_temp (id, user_id, date, cny, usd, exchange_rate, category, created_at, updated_at)
                            SELECT id, user_id, date, cny, usd, exchange_rate, category, created_at, updated_at
                            FROM investment_history
                        ''')
                        # 3. 删除原表
                        cursor.execute("DROP TABLE investment_history")
                        # 4. 重命名临时表
                        cursor.execute("ALTER TABLE investment_history_temp RENAME TO investment_history")
                conn.commit()
        except sqlite3.Error as e:
            print(f"创建表时发生错误: {e}")
            raise
    
    def add_record(self, user_id, date, cny, category='货币基金', deposit_start_date=None, deposit_duration=None, deposit_rate=None, usd=0, exchange_rate=0):
        """新增投资记录
        参数:
            user_id: 用户ID
            date: 日期，格式为YYYY-MM-DD
            cny: 人民币金额
            category: 投资类别，默认为'货币基金'
            deposit_start_date: 定存起始日期，格式为YYYY-MM-DD
            deposit_duration: 定存时长（月）
            deposit_rate: 定存年利率
            usd: 美元金额（保留字段）
            exchange_rate: 汇率（保留字段）
        返回:
            新增记录的ID
        """
        try:
            # 验证日期格式
            datetime.strptime(date, '%Y-%m-%d')
            
            # 如果是定存类别，验证定存相关字段
            if category == '定存':
                if deposit_start_date:
                    datetime.strptime(deposit_start_date, '%Y-%m-%d')
                if deposit_duration is not None and (not isinstance(deposit_duration, int) or deposit_duration <= 0):
                    raise ValueError("定存时长必须是正整数")
                if deposit_rate is not None and (not isinstance(deposit_rate, (int, float)) or deposit_rate < 0):
                    raise ValueError("年利率必须是非负数")
            
            with sqlite3.connect(self.db_path) as conn:
                cursor = conn.cursor()
                cursor.execute(
                    "INSERT INTO investment_history (user_id, date, cny, category, deposit_start_date, deposit_duration, deposit_rate, usd, exchange_rate) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)",
                    (user_id, date, cny, category, deposit_start_date, deposit_duration, deposit_rate, usd, exchange_rate)
                )
                conn.commit()
                return cursor.lastrowid
        except ValueError:
            raise ValueError("日期格式错误，请使用YYYY-MM-DD格式")
        except sqlite3.Error as e:
            print(f"新增记录时发生错误: {e}")
            raise
    
    def update_record(self, record_id, user_id, date=None, cny=None, category=None, deposit_start_date=None, deposit_duration=None, deposit_rate=None, usd=None, exchange_rate=None):
        """修改指定用户的投资记录
        参数:
            record_id: 记录ID
            user_id: 用户ID
            date: 日期，格式为YYYY-MM-DD
            cny: 人民币金额
            category: 投资类别
            deposit_start_date: 定存起始日期，格式为YYYY-MM-DD
            deposit_duration: 定存时长（月）
            deposit_rate: 定存年利率
            usd: 美元金额（保留字段）
            exchange_rate: 汇率（保留字段）
        返回:
            是否修改成功
        """
        try:
            fields = []
            params = []
            
            if date:
                # 验证日期格式
                datetime.strptime(date, '%Y-%m-%d')
                fields.append("date = ?")
                params.append(date)
                
            if cny is not None:
                fields.append("cny = ?")
                params.append(cny)
                
            if usd is not None:
                fields.append("usd = ?")
                params.append(usd)
                
            if exchange_rate is not None:
                fields.append("exchange_rate = ?")
                params.append(exchange_rate)
                
            if category is not None:
                fields.append("category = ?")
                params.append(category)
                
            if deposit_start_date is not None:
                fields.append("deposit_start_date = ?")
                params.append(deposit_start_date)
            
            if deposit_duration is not None:
                fields.append("deposit_duration = ?")
                params.append(deposit_duration)
            
            if deposit_rate is not None:
                fields.append("deposit_rate = ?")
                params.append(deposit_rate)
                
            # 添加更新时间
            fields.append("updated_at = ?")
            params.append(datetime.now().strftime('%Y-%m-%d %H:%M:%S'))
            
            # 如果是定存类别且有相关字段，验证定存相关字段
            if category == '定存' or (fields and 'category = ?' in fields and params[fields.index('category = ?')] == '定存'):
                for i, field in enumerate(fields):
                    if field == 'deposit_start_date = ?' and params[i]:
                        datetime.strptime(params[i], '%Y-%m-%d')
                    elif field == 'deposit_duration = ?' and params[i] is not None:
                        if not isinstance(params[i], int) or params[i] <= 0:
                            raise ValueError("定存时长必须是正整数")
                    elif field == 'deposit_rate = ?' and params[i] is not None:
                        if not isinstance(params[i], (int, float)) or params[i] < 0:
                            raise ValueError("年利率必须是非负数")
            
            # 添加记录ID
            params.append(record_id)
            
            if not fields:
                return False
            
            with sqlite3.connect(self.db_path) as conn:
                cursor = conn.cursor()
                sql = f"UPDATE investment_history SET {', '.join(fields)} WHERE id = ? AND user_id = ?"
                params.append(user_id)  # 添加user_id到参数列表末尾
                cursor.execute(sql, params)
                conn.commit()
                return cursor.rowcount > 0
        except ValueError:
            raise ValueError("日期格式错误，请使用YYYY-MM-DD格式")
        except sqlite3.Error as e:
            print(f"修改记录时发生错误: {e}")
            raise
    
    def delete_record(self, record_id, user_id):
        """删除指定用户的投资记录
        参数:
            record_id: 记录ID
            user_id: 用户ID
        返回:
            是否删除成功
        """
        try:
            with sqlite3.connect(self.db_path) as conn:
                cursor = conn.cursor()
                cursor.execute("DELETE FROM investment_history WHERE id = ? AND user_id = ?", (record_id, user_id))
                conn.commit()
                return cursor.rowcount > 0
        except sqlite3.Error as e:
            print(f"删除记录时发生错误: {e}")
            raise
    
    def get_record(self, record_id, user_id):
        """获取指定用户的单条投资记录
        参数:
            record_id: 记录ID
            user_id: 用户ID
        返回:
            记录字典，如果不存在或不属于该用户返回None
        """
        try:
            with sqlite3.connect(self.db_path) as conn:
                conn.row_factory = sqlite3.Row  # 使返回结果可以通过列名访问
                cursor = conn.cursor()
                cursor.execute("SELECT * FROM investment_history WHERE id = ? AND user_id = ?", (record_id, user_id))
                row = cursor.fetchone()
                return dict(row) if row else None
        except sqlite3.Error as e:
            print(f"获取记录时发生错误: {e}")
            raise
    
    def get_all_records(self, user_id, order_by='date DESC'):
        """获取指定用户的所有投资记录
        参数:
            user_id: 用户ID
            order_by: 排序方式，默认为按日期降序
        返回:
            记录列表
        """
        try:
            with sqlite3.connect(self.db_path) as conn:
                conn.row_factory = sqlite3.Row
                cursor = conn.cursor()
                cursor.execute(f"SELECT * FROM investment_history WHERE user_id = ? ORDER BY {order_by}", (user_id,))
                rows = cursor.fetchall()
                return [dict(row) for row in rows]
        except sqlite3.Error as e:
            print(f"获取所有记录时发生错误: {e}")
            raise
    
    def get_records_by_date_range(self, user_id, start_date, end_date):
        """获取指定用户在指定日期范围内的投资记录
        参数:
            user_id: 用户ID
            start_date: 开始日期，格式为YYYY-MM-DD
            end_date: 结束日期，格式为YYYY-MM-DD
        返回:
            记录列表
        """
        try:
            # 验证日期格式
            datetime.strptime(start_date, '%Y-%m-%d')
            datetime.strptime(end_date, '%Y-%m-%d')
            
            with sqlite3.connect(self.db_path) as conn:
                conn.row_factory = sqlite3.Row
                cursor = conn.cursor()
                cursor.execute(
                    "SELECT * FROM investment_history WHERE user_id = ? AND date BETWEEN ? AND ? ORDER BY date",
                    (user_id, start_date, end_date)
                )
                rows = cursor.fetchall()
                return [dict(row) for row in rows]
        except ValueError:
            raise ValueError("日期格式错误，请使用YYYY-MM-DD格式")
        except sqlite3.Error as e:
            print(f"获取日期范围内记录时发生错误: {e}")
            raise